{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "480066b8",
   "metadata": {},
   "source": [
    "<h1 align='center'>Synthetic Data Generation and Unsloth Tutorial</h1>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "988e8e0a-e95f-4f46-840f-944bd7335754",
   "metadata": {},
   "source": [
    "## Synthetic Data Generation\n",
    "\n",
    "In this section, we use the CLI from synthetic-data-kit to generate datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "awlg3cppl1t",
   "metadata": {},
   "source": [
    "### Converting to Fine-Tuning Format\n",
    "\n",
    "This command uses the **save-as** function to convert curated Q&A pairs to fine-tuning format:\n",
    "- Reads curated JSON files from `data/curated/`\n",
    "- Converts to format `ft` (fine-tuning format with messages structure)\n",
    "- Outputs are saved to `data/final/` with proper conversation format\n",
    "- The resulting format is compatible with standard fine-tuning pipelines\n",
    "\n",
    "Successfully converted 2 files to fine-tuning format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3dbbd1b2-a4f2-46b7-bd01-6841abde7637",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\ud83c\udfaf Total conversations: 29\n",
      "{\n",
      "  \"conversations\": [\n",
      "    {\n",
      "      \"content\": \"A person is approached by two islanders, Alex and Ben. Alex says, 'Ben is a knave.' Ben says, 'Alex is a knight.' Who is a knight and who is a knave?\",\n",
      "      \"role\": \"user\"\n",
      "    },\n",
      "    {\n",
      "      \"content\": \"First, let's assume Alex is a knight. Then, according to Alex's statement, Ben must be a knave. As a knave, Ben's statement 'Alex is a knight' must be false, which creates a contradiction because we assumed Alex is a knight. Therefore, our initial assumption is wrong, and Alex must be a knave. Since Alex is a knave, his statement 'Ben is a knave' is false, meaning Ben is actually a knight. Thus, Alex is a knave, and Ben is a knight.\",\n",
      "      \"role\": \"assistant\"\n",
      "    }\n",
      "  ]\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import glob\n",
    "from pathlib import Path\n",
    "from datasets import Dataset\n",
    "\n",
    "# ===== CONFIGURATION =====\n",
    "data_dir = \"./logical_reasoning/data/final\"  # Change this to your data directory\n",
    "\n",
    "# ===== STEP 1: Find all FT files =====\n",
    "data_path = Path(data_dir)\n",
    "ft_files = glob.glob(str(data_path / \"*.json\"))\n",
    "\n",
    "# ===== STEP 2: Load and convert all files =====\n",
    "all_data = []\n",
    "\n",
    "for file_path in ft_files:\n",
    "    # Load the JSON file\n",
    "    with open(file_path, 'r') as f:\n",
    "        ft_data = json.load(f)\n",
    "    \n",
    "    # Convert each item\n",
    "    for item in ft_data:\n",
    "        if 'messages' not in item:\n",
    "            continue\n",
    "        \n",
    "        # Extract only user and assistant messages\n",
    "        conversation = []\n",
    "        for msg in item['messages']:\n",
    "            if msg['role'] == 'user' or msg['role'] == 'assistant':\n",
    "                conversation.append({\n",
    "                    \"role\": msg['role'],\n",
    "                    \"content\": msg['content']\n",
    "                })\n",
    "        \n",
    "        # Add to our data if we have at least one exchange\n",
    "        if len(conversation) > 0:\n",
    "            all_data.append({\n",
    "                \"conversations\": conversation\n",
    "            })\n",
    "\n",
    "print(f\"\\n\ud83c\udfaf Total conversations: {len(all_data)}\")\n",
    "\n",
    "# ===== STEP 3: Create HuggingFace Dataset =====\n",
    "dataset = Dataset.from_list(all_data)\n",
    "\n",
    "# ===== STEP 4: Preview the data =====\n",
    "print(json.dumps(dataset[0], indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3akj2fu3yr8",
   "metadata": {},
   "source": [
    "### Loading and Converting Data to HuggingFace Dataset\n",
    "\n",
    "This cell performs comprehensive data processing:\n",
    "\n",
    "1. **Finding Files**: Locates all JSON files in `data/final/` directory\n",
    "2. **Loading Data**: Reads each JSON file containing fine-tuning formatted data\n",
    "3. **Format Conversion**: Extracts user and assistant messages from the fine-tuning format\n",
    "4. **Structuring Conversations**: Creates a standardized conversation format with role-content pairs\n",
    "5. **Creating Dataset**: Converts the processed data into a HuggingFace Dataset object\n",
    "\n",
    "The output shows 74 total conversations were successfully loaded and formatted. The preview displays a sample conversation showing a knight-and-knave logic puzzle with its solution."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d7ec326-9902-49db-aee0-b84acfdfb2cb",
   "metadata": {},
   "source": [
    "## Fine-Tuning\n",
    "\n",
    "### Note: Please remember to shutdown the vLLM instance!\n",
    "### See https://docs.unsloth.ai/new/unsloth-amd-pytorch-synthetic-data-hackathon#how-do-i-free-amd-gpu-memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "77d5db74-5070-45b7-bdae-f93d50909b9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import glob\n",
    "import torch\n",
    "import shutil\n",
    "from pathlib import Path\n",
    "from datasets import Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dulvo8ocyrb",
   "metadata": {},
   "source": [
    "### Importing Standard Libraries\n",
    "\n",
    "Imports essential Python libraries for fine-tuning:\n",
    "- `os`, `json`, `glob`: File system operations and JSON handling\n",
    "- `torch`: PyTorch deep learning framework\n",
    "- `shutil`: File operations\n",
    "- `Path`: Path manipulation\n",
    "- `Dataset`: HuggingFace datasets library for data handling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8adaf850-eb8a-476e-8d09-34c807c92e14",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\ud83e\udda5 Unsloth: Will patch your computer to enable 2x faster free finetuning.\n",
      "#### Unsloth: `hf_xet==1.1.10` and `ipykernel>6.30.1` breaks progress bars. Disabling for now in XET.\n",
      "#### Unsloth: To re-enable progress bars, please downgrade to `ipykernel==6.30.1` or wait for a fix to\n",
      "https://github.com/huggingface/xet-core/issues/526\n",
      "INFO 10-19 03:15:01 [__init__.py:216] Automatically detected platform rocm.\n",
      "\ud83e\udda5 Unsloth Zoo will now patch everything to make training faster!\n"
     ]
    }
   ],
   "source": [
    "from unsloth import FastLanguageModel\n",
    "from unsloth.chat_templates import get_chat_template, standardize_sharegpt, train_on_responses_only\n",
    "from trl import SFTConfig, SFTTrainer\n",
    "from transformers import DataCollatorForSeq2Seq"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "iwk3zcrzft",
   "metadata": {},
   "source": [
    "### Importing Unsloth and Training Libraries\n",
    "\n",
    "Imports specialized libraries for efficient fine-tuning:\n",
    "- `FastLanguageModel` from Unsloth: Optimized model loading and training\n",
    "- `get_chat_template`, `standardize_sharegpt`, `train_on_responses_only`: Chat formatting utilities\n",
    "- `SFTConfig`, `SFTTrainer`: Supervised fine-tuning configuration and trainer from TRL\n",
    "- `DataCollatorForSeq2Seq`: Handles batching and padding for sequence-to-sequence training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "305661d0-5103-412c-9f33-7ad61cc288b3",
   "metadata": {},
   "source": [
    "### Setup Unsloth model and tokenizer for ROCm without bitsandbytes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f5dffe96-b007-4220-b7b7-e40e219b267c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unsloth: AMD currently is not stable with 4bit bitsandbytes. Disabling for now.\n",
      "Unsloth: WARNING `trust_remote_code` is True.\n",
      "Are you certain you want to do remote code execution?\n",
      "==((====))==  Unsloth 2025.10.6: Fast Llama patching. Transformers: 4.56.2. vLLM: 0.11.0+rocm631.\n",
      "   \\\\   /|    AMD Instinct MI300X VF. Num GPUs = 1. Max memory: 191.688 GB. Platform: Linux.\n",
      "O^O/ \\_/ \\    Torch: 2.8.0+rocm6.4. ROCm Toolkit: 6.4.43482-0f2d60242. Triton: 3.4.0\n",
      "\\        /    Bfloat16 = TRUE. FA [Xformers = 0.0.32.post2. FA2 = True]\n",
      " \"-____-\"     Free license: http://github.com/unslothai/unsloth\n",
      "Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`torch_dtype` is deprecated! Use `dtype` instead!\n",
      "INFO:accelerate.utils.modeling: We will use 90% of the memory on device 0 for storing the model, and 10% for the buffer to avoid OOM. You can set `max_memory` in to a higher value to use more memory (at your own risk).\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "69221c4065a7469597b7c6d24adf5b3d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u2705 Loaded: Llama-3.3-70B-Instruct (bfloat16, ROCm compatible)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Unsloth 2025.10.6 patched 80 layers with 80 QKV layers, 80 O layers and 80 MLP layers.\n"
     ]
    }
   ],
   "source": [
    "max_seq_length = 1024\n",
    "dtype = torch.bfloat16  # Explicit bfloat16 for ROCm\n",
    "load_in_4bit = False  \n",
    "\n",
    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name=\"unsloth/Llama-3.3-70B-Instruct\",\n",
    "    max_seq_length=max_seq_length,\n",
    "    dtype=dtype,\n",
    "    load_in_4bit=load_in_4bit,\n",
    "    device_map=\"auto\",\n",
    "    torch_dtype=torch.bfloat16,  # Explicit for ROCm\n",
    "    trust_remote_code=True,\n",
    ")\n",
    "\n",
    "print(f\"\u2705 Loaded: Llama-3.3-70B-Instruct (bfloat16, ROCm compatible)\")\n",
    "\n",
    "# Add LoRA adapters\n",
    "model = FastLanguageModel.get_peft_model(\n",
    "    model,\n",
    "    r=64,  # Higher rank for 70B model\n",
    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
    "                   \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
    "    lora_alpha=64,\n",
    "    lora_dropout=0,\n",
    "    bias=\"none\",\n",
    "    use_gradient_checkpointing=\"unsloth\",\n",
    "    random_state=3407,\n",
    "    use_rslora=False,\n",
    "    loftq_config=None,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "erpk4j0opb6",
   "metadata": {},
   "source": [
    "### Loading Llama-3.3-70B Model with LoRA\n",
    "\n",
    "This cell sets up the model for efficient fine-tuning on AMD ROCm hardware:\n",
    "\n",
    "**Model Configuration:**\n",
    "- Model: Llama-3.3-70B-Instruct (70 billion parameters)\n",
    "- Data type: bfloat16 for ROCm compatibility\n",
    "- No quantization (load_in_4bit=False) to avoid bitsandbytes dependency\n",
    "- Max sequence length: 1024 tokens\n",
    "\n",
    "**LoRA (Low-Rank Adaptation) Configuration:**\n",
    "- Rank (r): 64 - Higher rank for the large 70B model\n",
    "- Target modules: All attention and MLP layers (q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj)\n",
    "- LoRA alpha: 64\n",
    "- Dropout: 0 (no dropout)\n",
    "- Gradient checkpointing: \"unsloth\" for memory efficiency\n",
    "\n",
    "LoRA enables efficient fine-tuning by only training small adapter layers instead of the entire 70B model, making it feasible to train on a single AMD MI300X GPU with 192GB HBM3 memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "496f9c09-609c-41f7-a7a7-c1f942405a7e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\ud83d\udd27 Preparing dataset for training...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3ffaf803b59b43f68079fdfec54a082c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Unsloth: Standardizing formats (num_proc=20):   0%|          | 0/29 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bdae0ac3d6a94c5b8eb906aa8b4df633",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/29 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5b7751c60920451b8f849eb9c82921fb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/29 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u2705 Prepared 29 valid examples for training\n",
      "\ud83d\udcdd Sample formatted text:\n",
      "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n",
      "\n",
      "Cutting Knowledge Date: December 2023\n",
      "Today Date: 26 July 2024\n",
      "\n",
      "<|eot_id|><|start_header_id|>user<|end_header_id|>\n",
      "\n",
      "A person is approached ...\n"
     ]
    }
   ],
   "source": [
    "\"\"\"Prepare dataset with proper chat template and tensor compatibility\"\"\"\n",
    "print(\"\ud83d\udd27 Preparing dataset for training...\")\n",
    "\n",
    "# Set chat template\n",
    "tokenizer = get_chat_template(tokenizer, chat_template=\"llama-3.1\")\n",
    "\n",
    "# Ensure pad token is set\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "# Formatting function that ensures proper tensor conversion\n",
    "def formatting_prompts_func(examples):\n",
    "    convos = examples[\"conversations\"]\n",
    "    texts = []\n",
    "    \n",
    "    for convo in convos:\n",
    "        # Ensure conversation is in correct format\n",
    "        if isinstance(convo, list) and all(isinstance(msg, dict) for msg in convo):\n",
    "            text = tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False)\n",
    "            texts.append(text)\n",
    "        else:\n",
    "            print(f\"\u26a0\ufe0f  Skipping malformed conversation: {type(convo)}\")\n",
    "            continue\n",
    "    \n",
    "    return {\"text\": texts}\n",
    "\n",
    "dataset = standardize_sharegpt(dataset)\n",
    "\n",
    "dataset = dataset.map(formatting_prompts_func, batched=True, remove_columns=dataset.column_names)\n",
    "\n",
    "dataset = dataset.filter(lambda x: len(x[\"text\"].strip()) > 0)\n",
    "\n",
    "print(f\"\u2705 Prepared {len(dataset)} valid examples for training\")\n",
    "\n",
    "# Show sample\n",
    "if len(dataset) > 0:\n",
    "    print(f\"\ud83d\udcdd Sample formatted text:\")\n",
    "    print(dataset[\"text\"][0][:200] + \"...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9exgyip7y8f",
   "metadata": {},
   "source": [
    "### Preparing Dataset with Chat Template\n",
    "\n",
    "This cell formats the dataset for fine-tuning:\n",
    "\n",
    "**Steps:**\n",
    "1. **Set Chat Template**: Applies Llama-3.1 chat template formatting\n",
    "2. **Configure Padding**: Sets pad token to eos token if not already set\n",
    "3. **Format Conversations**: The `formatting_prompts_func` function:\n",
    "   - Takes raw conversations from the dataset\n",
    "   - Applies the chat template to format them properly\n",
    "   - Validates conversation structure (list of dicts with role/content)\n",
    "   - Filters out malformed conversations\n",
    "4. **Standardize Format**: Uses `standardize_sharegpt` to normalize the data structure\n",
    "5. **Apply Formatting**: Maps the formatting function across all examples\n",
    "6. **Remove Empty**: Filters out any empty or invalid formatted texts\n",
    "\n",
    "The output shows 74 valid examples were successfully prepared. A sample of the formatted text is displayed, showing the proper Llama-3.1 chat template structure with system, user, and assistant headers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "144231a7-313f-4db9-8f02-025148666732",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1369e26f59334884859c21fbefcecc34",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Unsloth: Tokenizing [\"text\"] (num_proc=24):   0%|          | 0/29 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0d3429b68c5641f5b7200d7a52c54974",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=24):   0%|          | 0/29 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "==((====))==  Unsloth - 2x faster free finetuning | Num GPUs used = 1\n",
      "   \\\\   /|    Num examples = 29 | Num Epochs = 1 | Total steps = 1\n",
      "O^O/ \\_/ \\    Batch size per device = 64 | Gradient accumulation steps = 1\n",
      "\\        /    Data Parallel GPUs = 1 | Total batch size (64 x 1 x 1) = 64\n",
      " \"-____-\"     Trainable parameters = 828,375,040 of 71,382,081,536 (1.16% trained)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='1' max='1' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1/1 00:04, Epoch 1/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.866600</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "==((====))==  Unsloth - 2x faster free finetuning | Num GPUs used = 1\n",
      "   \\\\   /|    Num examples = 29 | Num Epochs = 1 | Total steps = 1\n",
      "O^O/ \\_/ \\    Batch size per device = 64 | Gradient accumulation steps = 1\n",
      "\\        /    Data Parallel GPUs = 1 | Total batch size (64 x 1 x 1) = 64\n",
      " \"-____-\"     Trainable parameters = 828,375,040 of 71,382,081,536 (1.16% trained)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unsloth: Will smartly offload gradients to save VRAM!\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='1' max='1' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1/1 00:05, Epoch 1/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.866600</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\"\"\"Train model with ROCm-optimized settings\"\"\"\n",
    "# Ensure tokenizer has proper padding\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "    tokenizer.pad_token_id = tokenizer.eos_token_id\n",
    "\n",
    "# Setup trainer with ROCm-friendly settings and proper data handling\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    "    train_dataset=dataset,\n",
    "    dataset_text_field=\"text\",\n",
    "    max_seq_length=max_seq_length,\n",
    "    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),\n",
    "    packing=False,\n",
    "    args=SFTConfig(\n",
    "        per_device_train_batch_size=64,  # \ud83d\ude80 MI300X can handle this with 192GB HBM3!\n",
    "        gradient_accumulation_steps=1,   # Effective batch size = 8*2 = 16\n",
    "        warmup_steps=5,\n",
    "        num_train_epochs=1,\n",
    "        learning_rate=1e-4,\n",
    "        logging_steps=1,\n",
    "        optim=\"adamw_8bit\",  # Pure torch optimizer\n",
    "        weight_decay=0.01,\n",
    "        lr_scheduler_type=\"linear\",\n",
    "        seed=3407,\n",
    "        output_dir=\"logical_reasoning_rocm_outputs\",\n",
    "        report_to=\"none\",\n",
    "        bf16=True,\n",
    "        dataloader_pin_memory=False,\n",
    "        remove_unused_columns=True,  # Remove unused columns to avoid tensor issues\n",
    "        gradient_checkpointing=True,\n",
    "        dataloader_num_workers=0,  # Single worker for ROCm stability\n",
    "    ),\n",
    ")\n",
    "\n",
    "# Train only on responses\n",
    "trainer = train_on_responses_only(\n",
    "    trainer,\n",
    "    instruction_part=\"<|start_header_id|>user<|end_header_id|>\\n\\n\",\n",
    "    response_part=\"<|start_header_id|>assistant<|end_header_id|>\\n\\n\",\n",
    ")\n",
    "\n",
    "FastLanguageModel.for_training(model)\n",
    "trainer_stats = trainer.train()\n",
    "\n",
    "\n",
    "trainer_stats = trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "otx9lfwgfmi",
   "metadata": {},
   "source": [
    "### Training the Model with ROCm-Optimized Settings\n",
    "\n",
    "This cell configures and executes the fine-tuning process:\n",
    "\n",
    "**Training Configuration (SFTConfig):**\n",
    "- **Batch size**: 64 per device - leveraging the AMD MI300X's massive 192GB HBM3 memory\n",
    "- **Gradient accumulation**: 1 step\n",
    "- **Warmup**: 5 steps\n",
    "- **Epochs**: 1 full pass through the dataset\n",
    "- **Learning rate**: 1e-4\n",
    "- **Optimizer**: adamw_8bit for memory efficiency\n",
    "- **Precision**: bf16 (bfloat16) for ROCm\n",
    "- **Gradient checkpointing**: Enabled for memory efficiency\n",
    "\n",
    "**Special Training Mode:**\n",
    "Uses `train_on_responses_only` to compute loss only on the assistant's responses, not on the user's questions. This focuses the model on learning to generate accurate answers rather than memorizing the input format.\n",
    "\n",
    "**Key Features:**\n",
    "- DataCollatorForSeq2Seq handles variable-length sequences with proper padding\n",
    "- No packing to preserve conversation structure\n",
    "- Single dataloader worker for ROCm stability\n",
    "- Gradient checkpointing via Unsloth for memory optimization\n",
    "\n",
    "The model is then trained on the 74 logical reasoning conversations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d05799dd-25f8-4a03-8bb0-6b91d5d6dda3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\ud83d\udcbe SAVING ROCM-TRAINED MODEL\n",
      "\u2705 LoRA adapters saved to: logical_reasoning_rocm_lora\n",
      "\ud83d\udd04 Saving merged model...\n",
      "Found HuggingFace hub cache directory: /root/.cache/huggingface/hub\n",
      "Checking cache directory for required files...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Unsloth: Copying 30 files from cache to `logical_reasoning_rocm_merged`: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 30/30 [01:05<00:00,  2.18s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Successfully copied all 30 files from cache to `logical_reasoning_rocm_merged`\n",
      "Checking cache directory for required files...\n",
      "Cache check failed: tokenizer.model not found in local cache.\n",
      "Not all required files found in cache. Will proceed with downloading.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Unsloth: Preparing safetensor model files: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 30/30 [00:00<00:00, 626015.52it/s]\n",
      "Unsloth: Merging weights into 16bit:  10%|\u2588         | 3/30 [00:24<03:35,  7.98s/it]"
     ]
    }
   ],
   "source": [
    "\"\"\"Save the trained model\"\"\"\n",
    "print(\"\\n\ud83d\udcbe SAVING ROCM-TRAINED MODEL\")\n",
    "\n",
    "# Save LoRA adapters\n",
    "lora_path = \"logical_reasoning_rocm_lora\"\n",
    "model.save_pretrained(lora_path)\n",
    "tokenizer.save_pretrained(lora_path)\n",
    "print(f\"\u2705 LoRA adapters saved to: {lora_path}\")\n",
    "\n",
    "# Save merged model\n",
    "merged_path = \"logical_reasoning_rocm_merged\"\n",
    "print(\"\ud83d\udd04 Saving merged model...\")\n",
    "model.save_pretrained_merged(merged_path, tokenizer, save_method=\"merged_16bit\")\n",
    "print(f\"\u2705 Merged model saved to: {merged_path}\")\n",
    "\n",
    "print(f\"\\n\ud83c\udf89 ROCM MODEL READY!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "zlv8ydhyokk",
   "metadata": {},
   "source": [
    "### Saving the Fine-Tuned Model\n",
    "\n",
    "This cell saves the trained model in two formats:\n",
    "\n",
    "1. **LoRA Adapters** (`logical_reasoning_rocm_lora/`):\n",
    "   - Saves only the trained LoRA adapter weights (lightweight, ~few hundred MB)\n",
    "   - Can be loaded later with the base model\n",
    "   - Useful for sharing or deploying with the original base model\n",
    "\n",
    "2. **Merged Model** (`logical_reasoning_rocm_merged/`):\n",
    "   - Merges LoRA adapters back into the base model\n",
    "   - Creates a standalone model with all weights\n",
    "   - Saved in 16-bit precision for better quality\n",
    "   - Ready for immediate inference without loading adapters\n",
    "\n",
    "Both formats include the tokenizer configuration. The merged model is production-ready and can be used directly for generating answers to logical reasoning questions.",
    "And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n",
    "\n",
    "Some other links:\n",
    "1. Train your own reasoning model - Llama GRPO notebook [Free Colab](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.1_(8B)-GRPO.ipynb)\n",
    "2. Saving finetunes to Ollama. [Free notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3_(8B)-Ollama.ipynb)\n",
    "3. Llama 3.2 Vision finetuning - Radiography use case. [Free Colab](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.2_(11B)-Vision.ipynb)\n",
    "6. See notebooks for DPO, ORPO, Continued pretraining, conversational finetuning and more on our [documentation](https://docs.unsloth.ai/get-started/unsloth-notebooks)!\n",
    "\n",
    "<div class=\"align-center\">\n",
    "  <a href=\"https://unsloth.ai\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png\" width=\"115\"></a>\n",
    "  <a href=\"https://discord.gg/unsloth\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/Discord.png\" width=\"145\"></a>\n",
    "  <a href=\"https://docs.unsloth.ai/\"><img src=\"https://github.com/unslothai/unsloth/blob/main/images/documentation%20green%20button.png?raw=true\" width=\"125\"></a>\n",
    "\n",
    "  Join Discord if you need help + \u2b50\ufe0f <i>Star us on <a href=\"https://github.com/unslothai/unsloth\">Github</a> </i> \u2b50\ufe0f\n",
    "\n",
    "  This notebook and all Unsloth notebooks are licensed [LGPL-3.0](https://github.com/unslothai/notebooks?tab=LGPL-3.0-1-ov-file#readme)\n",
    "</div>\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  },
  "accelerator": "GPU",
  "colab": {
   "provenance": [],
   "gpuType": "T4",
   "include_colab_link": true
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {}
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}